# # * main training use

# # # * single machine single gpu
# ACC_CONFIG_FILE="configs/acc_single_default.yaml"
# GPU_IDS="0"
# NGPU=1
# MASTER_PORT=29500
# accelerate launch --config_file $ACC_CONFIG_FILE --main_process_port $MASTER_PORT --num_processes $NGPU --gpu_ids $GPU_IDS \
# scripts/train.py \
#     --config reflow/configs/train.py \
#     --workdir "logs/tmp" \
#     --comment ""

# # # * single machine multi gpu
# ACC_CONFIG_FILE="configs/acc_multi_default.yaml"
# GPU_IDS="0,1,2,3,4,5,6,7"
# NGPU=8
# MASTER_PORT=29500
# accelerate launch --config_file $ACC_CONFIG_FILE --main_process_port $MASTER_PORT --num_processes $NGPU --gpu_ids $GPU_IDS \
# scripts/train.py \
#     --config reflow/configs/train.py \
#     --workdir "logs/2reflow_add_cfg" \
#     --comment "2reflow training with lora ; p_uncond=0.1 ; total bs = (8gpu x 10bs x 1acc) = 80 ; 5M data , 62500 steps ; lr 1000 stes warmup then constant , peak 1e-4"


# # # # * multi machine training
# # ACC_CONFIG_FILE="configs/acc_multi_machine.yaml"
# # NUM_NODE=2
# # NGPU_PER_NODE=8
# # NUM_PROCESSES=`expr $NUM_NODE \* $NGPU_PER_NODE`
# # # RANK=3 # ! existing in the system already
# # # MASTER_PORT=29500 # maybe useful

# # # echo "rank $RANK machine started , $NUM_NODE machine in total" > /home/zhoupengfei886/tmp.txt
# # # echo $MASTER_ADDR >> /home/zhoupengfei886/tmp.txt

# # accelerate launch --config_file $ACC_CONFIG_FILE --num_machines $NUM_NODE --num_processes $NUM_PROCESSES --machine_rank $RANK --main_process_ip $MASTER_ADDR scripts/train.py \
# #     --config reflow/configs/train.py \
# #     --workdir logs/distill_1step_2reflow_v2 \
# #     --comment "2reflow training overfit 10k val data ; p_uncond=0.1 ; total bs = (1gpu x 2bs x 8acc) = 16 ; 10k data, 2epoch, 20000/16=1250 ; lr 0 stes warmup then constant , peak 1e-4 ; use lpips loss on latent"


# # * sample part

# # # 1. caption file (请自由准备)
# # # 2. 用脚本采样
# # # 3. 用脚本计算指标并保存结果

# # # 通常的使用 case
# # # 固定 captions, ckpt_path, seed, save_dir, infer_steps, guidance_scale, devices, bs
# # # 依次生成 IS, FID, CLIPSim 的 samples

# # # # * sample metric 数据
# devices="1,2,3,4,5,6,7"
# python scripts/metric_sample.py \
#     --metric_type FID \
#     --caption_path data/reflow/coco2017_val_random5k.txt \
#     --save_dir "samples/metric_sample/coco2017_val_5k/sdv1-4_laion_1MPrompts_2reflow_lpips+l2_N=25_sampleCFG=1.0" \
#     --devices $devices \
#     --infer_steps=25 \
#     --guidance_scale=1.0 \
#     --bs=10 \
#     --pipeline stable_diffusion \
#     --scheduler euler_dummy \
#     --pipeline_ckpt checkpoints/SD-1-4 \
#     --ckpt_path logs/online/sdv1-4_laion_1MPrompts_2reflow_lpips+l2/checkpoints/score_model_s40000.pth \
#     --use_xformers

# devices="2,3,4,5,6,7"
# N=25
# cap_list=(AutoCaption_10w_head_1 AutoCaption_10w_head_10 AutoCaption_10w_head_100 AutoCaption_10w_head_1k AutoCaption_10w_head_1w AutoCaption_10w_full laion_random10W)
# for i in "${!cap_list[@]}"
# do
#     caption_type=${cap_list[$i]}
#     python scripts/metric_sample.py \
#         --metric_type FID \
#         --caption_path data/reflow/coco2017_val_random5k.txt \
#         --save_dir "samples/metric_sample/ablation/$caption_type/N=$N" \
#         --devices $devices \
#         --infer_steps=$N \
#         --guidance_scale=2 \
#         --bs=10 \
#         --pipeline stable_diffusion \
#         --scheduler euler_dummy \
#         --pipeline_ckpt checkpoints/SD-1-4 \
#         --ckpt_path "logs/pokemon/2reflow/$caption_type/score_model_s2000.pth" \
#         --use_xformers
#     python scripts/metric_sample.py \
#         --metric_type CLIPSim \
#         --num_candidates 8 \
#         --caption_path data/reflow/coco2017_val_random5k.txt \
#         --save_dir "samples/metric_sample/ablation/$caption_type/N=$N" \
#         --devices $devices \
#         --infer_steps=$N \
#         --guidance_scale=2 \
#         --bs=10 \
#         --pipeline stable_diffusion \
#         --scheduler euler_dummy \
#         --pipeline_ckpt checkpoints/SD-1-4 \
#         --ckpt_path "logs/pokemon/2reflow/$caption_type/score_model_s2000.pth" \
#         --use_xformers
# done



# # * generate reflow data
# python scripts/generate_reflow_data.py \
#     --caption_path data/reflow/laion6+_random50k.txt \
#     --save_dir data/reflow/laion_random50k \
#     --devices "0,1,2,3" \
#     --infer_steps 25 \
#     --seed 2891 \
#     --scheduler dpm_solver_multi \
#     --pipeline_ckpt checkpoints/AltDiffusion \
#     --use_xformers


# # * generate coco original data

# caption_path="tmp/coco2014_train_random1M.txt"
# devices="0,1,2,3"
# split=train
# bs=50
# dl_workers=4

# python scripts/generate_coco_data_retrieve_noise.py \
#     --caption_path $caption_path \
#     --save_dir "data/coco2014_reflow/coco_latent_train1M" \
#     --devices $devices \
#     --split $split \
#     --bs $bs \
#     --dl_workers $dl_workers 

# python scripts/generate_coco_data_retrieve_noise.py \
#     --caption_path $caption_path \
#     --save_dir "data/coco2014_reflow/coco_latent_train1M-retrieve_alt_gen_train5M_noise" \
#     --devices $devices \
#     --split $split \
#     --bs $bs \
#     --dl_workers $dl_workers \
#     --retrive_noise \
#     --reflow_ds_path "data/coco2014_reflow/alt_gen_train5M"

# # * sample reflow ds
# python scripts/sample_reflow_ds.py \
#     --config reflow/configs/sample.py \
#     --eval_folder samples/ds_sample/euler25/2reflow_v3_s125000 \
#     --config.diffusers.score_model_ckpt logs/2reflow_AltInit_v3/checkpoints/score_model_s125000.pth \
#     --config.data.eval_root data/coco2014_reflow/alt_gen_train_first100 \
#     --config.device 'cuda:0'

# # * train online
# # ! finetune sd 的话也是指定 reflow ckpt inference/train 
ACC_CONFIG_FILE="configs/acc_multi_default.yaml"
GPU_IDS="0,1,2,3,4,5,6,7"
NGPU=8
MASTER_PORT=29501
accelerate launch --config_file $ACC_CONFIG_FILE --main_process_port $MASTER_PORT --num_processes $NGPU --gpu_ids $GPU_IDS \
scripts/train_online.py \
    --config reflow/configs/train_online.py \
    --workdir "logs/coco_laion_full/2reflow" \
    --config.diffusers.scheduler "dpm_solver_multi" \
    --config.diffusers.reflow_ckpt_inference "" \
    --config.diffusers.reflow_ckpt_train "logs/online/sdv1-4_laion_1MPrompts_guidance7.5/checkpoints/score_model_s40000.pth" \
    --config.diffusers.num_inference_steps=25 \
    --config.diffusers.guidance_scale=7.5 \
    --config.training.num_steps=20000 \
    --config.training.batch_size=32 \
    --config.training.gradient_accumulation_steps=1 \
    --config.training.ckpt_path "" \
    --config.training.p_uncond=0.1 \
    --config.training.snapshot_freq=2000 \
    --config.reflow.reflow_t_schedule uniform \
    --config.reflow.reflow_loss l2 \
    --config.data.caption_path "data/reflow/coco_laion_full.txt" \
    --config.seed=58239057 \
    --config.optim.lr 1e-5 \
    --comment "" \
    
    # --config.reflow.finetune_sd "yes" \
    # --config.reflow.sd_t_schedule "t0"



# # # * train lora
# ACC_CONFIG_FILE="configs/acc_multi_default.yaml"
# GPU_IDS="4,5,6,7"
# NGPU=4
# MASTER_PORT=29502
# accelerate launch --config_file $ACC_CONFIG_FILE --main_process_port $MASTER_PORT --num_processes $NGPU --gpu_ids $GPU_IDS \
# scripts/lora/train_lora_1reflow.py \
#     --pretrained_model_name_or_path checkpoints/SD-1-4 \
#     --dataset_name "data/pokemon/data" \
#     --mixed_precision fp16 \
#     --dataloader_num_workers=4 \
#     --resolution=512 --center_crop --random_flip \
#     --train_batch_size=8 \
#     --gradient_accumulation_steps=1 \
#     --max_train_steps=2000 \
#     --learning_rate=1e-4 \
#     --adam_weight_decay=0 \
#     --max_grad_norm=1 \
#     --lr_scheduler="constant" --lr_warmup_steps=0 \
#     --output_dir logs/lora/lr1e-4_rank4_s2000 \
#     --report_to tensorboard \
#     --checkpointing_steps=2000 \
#     --validation_epochs=10000 \
#     --validation_prompt "a drawing of a green pokemon with red eyes" \
#     --seed=241355 \
#     --lora_rank=4 \
#     --reflow_ckpt_path "logs/online/sdv1-4_laion_1MPrompts_guidance7.5/checkpoints/score_model_s40000.pth"


# # # * train 1reflow
# ACC_CONFIG_FILE="configs/acc_multi_default.yaml"
# GPU_IDS="0,1,2,3,4,5,6,7"
# NGPU=8
# MASTER_PORT=29501
# ds_list=(emoji nouns sketch-scene pokemon)
# for ds_name in ${ds_list[*]}
# do
#     accelerate launch --config_file $ACC_CONFIG_FILE --main_process_port $MASTER_PORT --num_processes $NGPU --gpu_ids $GPU_IDS \
#     scripts/train_1reflow.py \
#         --pretrained_model_name_or_path checkpoints/SD-1-4 \
#         --dataset_names "data/$ds_name/data" \
#         --mixed_precision fp16 \
#         --dataloader_num_workers=8 \
#         --resolution=512 --center_crop --random_flip \
#         --train_batch_size=32 \
#         --gradient_accumulation_steps=1 \
#         --max_train_steps=1000 \
#         --max_grad_norm=1 \
#         --learning_rate=1e-5 \
#         --lr_scheduler="constant_with_warmup" --lr_warmup_steps=0 \
#         --adam_weight_decay=0 \
#         --output_dir "logs/$ds_name/finetune_sd/main" \
#         --report_to tensorboard \
#         --checkpointing_steps=500 \
#         --validation_epochs=10000 \
#         --seed=567823 \
#         --reflow_ckpt_path "" \
#         --reflow_loss "l2" \
#         --gradient_checkpointing \
#         --enable_xformers_memory_efficient_attention \
#         --finetune_sd
# done


# # # * train online
# # # ! finetune sd 的话也是指定 reflow ckpt inference/train
# ACC_CONFIG_FILE="configs/acc_multi_default.yaml"
# GPU_IDS="0,1,2,3,4,5,6,7"
# NGPU=8
# MASTER_PORT=29501
# accelerate launch --config_file $ACC_CONFIG_FILE --main_process_port $MASTER_PORT --num_processes $NGPU --gpu_ids $GPU_IDS \
# scripts/train_online.py \
#     --config reflow/configs/train_online.py \
#     --workdir "logs/nouns/distill_sd/main" \
#     --config.diffusers.scheduler "dpm_solver_multi" \
#     --config.diffusers.reflow_ckpt_inference "logs/nouns/finetune_sd/main/score_model_s1000.pth" \
#     --config.diffusers.reflow_ckpt_train "logs/nouns/finetune_sd/main/score_model_s1000.pth" \
#     --config.diffusers.num_inference_steps=25 \
#     --config.diffusers.guidance_scale=7.5 \
#     --config.training.num_steps=2000 \
#     --config.training.batch_size=8 \
#     --config.training.gradient_accumulation_steps=4 \
#     --config.training.ckpt_path "" \
#     --config.training.p_uncond=0 \
#     --config.training.snapshot_freq=1000 \
#     --config.reflow.reflow_t_schedule t0 \
#     --config.reflow.reflow_loss l2 \
#     --config.data.caption_path "data/reflow/laion_auto_caption_1W.txt" \
#     --config.seed=249057 \
#     --config.optim.lr 1e-5 \
#     --comment "" \
#     --config.reflow.finetune_sd "yes" \
#     --config.reflow.sd_t_schedule "t0"